import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt


# 1、生成一些随机数据
n_samples = 100
data = torch.randn(n_samples, 2)  # 随机生成输入数据

# 点在圆内为1，点在圆外为0
labels = (data[:, 0] ** 2 + data[:, 1] ** 2 < 1).float().unsqueeze(1)  # 标签为1或0

print(data)  # 打印数据
print(labels)  # 打印标签

# 可视化数据
# plt.scatter(data[:, 0], data[:, 1], c=labels.squeeze(), cmap='coolwarm', edgecolor='k')
# plt.title('Generated Data')
# plt.xlabel('F1')
# plt.ylabel('F2')
# plt.show()

# 2、创建一个前馈神经网络模型
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(2, 4)  # 输入层到隐藏层
        self.fc2 = nn.Linear(4, 1)  # 隐藏层到输出层
        self.sigmoid = nn.Sigmoid()  # Sigmoid激活函数

    def forward(self, x):
        x = torch.relu(self.fc1(x))  # 激活函数
        x = torch.sigmoid(self.fc2(x))  # 输出层激活函数
        return x

model = SimpleNN()

# 3、定义损失函数和优化器
criterion = nn.BCELoss()  # 二分类交叉熵损失函数    
optimizer = optim.SGD(model.parameters(), lr=0.01)  # 随机梯度下降优化器

# 4、训练模型
epochs = 100000
for epoch in range(epochs):
    outputs = model(data)  # 前向传播
    loss = criterion(outputs, labels)  # 计算损失
    optimizer.zero_grad()  # 清空梯度
    loss.backward()  # 反向传播
    optimizer.step()  # 更新参数
    if (epoch + 1) % 10 == 0: # 每 10 轮打印一次损失
        print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.4f}')  # 打印损失

# 5、可视化训练结果
def plot_decision_boundary(model, data):
    x_min, x_max = data[:, 0].min() - 1, data[:, 0].max() + 1
    y_min, y_max = data[:, 1].min() - 1, data[:, 1].max() + 1
    xx, yy = torch.meshgrid(torch.arange(x_min, x_max, 0.1), torch.arange(y_min, y_max, 0.1), indexing='ij')
    grid = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1)], dim=1)
    predictions = model(grid).detach().numpy().reshape(xx.shape)
    plt.contourf(xx, yy, predictions, levels=[0, 0.5, 1], cmap='coolwarm', alpha=0.7)
    plt.scatter(data[:, 0], data[:, 1], c=labels.squeeze(), cmap='coolwarm', edgecolors='k')
    plt.title("Decision Boundary")
    plt.savefig("decision_boundary.png")

plot_decision_boundary(model, data)